import keras
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler


def load_data():
    cal_housing = np.loadtxt(
        "C:/Users/MSI/scikit_learn_data/CaliforniaHousing/cal_housing.data", delimiter=","
    )
    # Columns are not in the same order compared to the previous
    # URL resource on lib.stat.cmu.edu
    columns_index = [8, 7, 2, 3, 4, 5, 6, 1, 0]
    cal_housing = cal_housing[:, columns_index]
    target, data = cal_housing[:, 0], cal_housing[:, 1:]

    # avg rooms = total rooms / households
    data[:, 2] /= data[:, 5]

    # avg bed rooms = total bed rooms / households
    data[:, 3] /= data[:, 5]

    # avg occupancy = population / households
    data[:, 5] = data[:, 4] / data[:, 5]

    # target in units of 100,000
    target = target / 100000.0
    return data, target

data, target = load_data()
print(data.shape)
print(target.shape)

# 切割数据
x_train_all, x_test, y_train_all, y_test = train_test_split(data, target, random_state=1)
x_train, x_valid, y_train, y_valid = train_test_split(x_train_all, y_train_all, random_state=1)
print(x_train.shape, x_valid.shape, y_train.shape, y_valid.shape)
print(x_test.shape, y_test.shape)

# 标准化处理
scaler = StandardScaler()
x_train_scaled = scaler.fit_transform(x_train)
x_valid_scaled = scaler.transform(x_valid)


# 子类api
class WideDeepModel(keras.models.Model):
    def __init__(self):
        """定义模型的层次"""
        super().__init__()
        self.hidden1 = keras.layers.Dense(32, activation="relu")
        self.hidden2 = keras.layers.Dense(32, activation="relu")
        self.output_layer = keras.layers.Dense(1)

    def call(self, input):
        hidden1 = self.hidden1(input)
        hidden2 = self.hidden2(hidden1)
        # 拼接
        concat = keras.layers.concatenate([input, hidden2])
        output_layer = self.output_layer(concat)
        return output_layer


model = WideDeepModel()
model.build(input_shape=(None, 8))
model.compile(loss="mse", optimizer="adam", metrics=["mse"])
model.fit(x_train_scaled, y_train, validation_data=(x_valid_scaled, y_valid), epochs=20)
